﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Linq.Expressions;
using System.Collections;
using CoolCode.Linq.Expressions;

namespace CoolCode.Linq
{
    public static class ExpandableQueryExtensions
    {
        /// <summary>
        /// Refer to http://www.albahari.com/nutshell/linqkit.html and
        /// http://tomasp.net/blog/linq-expand.aspx for more information.
        /// </summary>
        public static IQueryable<T> AsExpandable<T>(this IQueryable<T> query)
        {
            return query.WithProvider<T, ExpandableQueryProvider>();
        }

        public static IQueryable<T> WithProvider<T, TProvider>(this IQueryable<T> query)
            where TProvider : AbstractQueryProvider
        {
            if (query is ExpandableQuery<T, TProvider>)
                return query;

            return new ExpandableQuery<T, TProvider>(query);
        }

        public static IQueryable<T> WithProvider<T>(this IQueryable<T> query, Type TProvider)
        {
            Type type = typeof(ExpandableQuery<,>).MakeGenericType(typeof(T), TProvider);

            if (type.IsInstanceOfType(query))
                return query;

            return (IQueryable<T>)Activator.CreateInstance(type, query);
        }
    }

    /// <summary>
    /// An IQueryable wrapper that allows us to visit the query's expression tree just before LINQ to SQL gets to it.
    /// This is based on the excellent work of Tomas Petricek: http://tomasp.net/blog/linq-expand.aspx
    /// </summary>
    public class ExpandableQuery<T, TProvider> : IQueryable<T>, IOrderedQueryable<T>, IOrderedQueryable
        where TProvider : AbstractQueryProvider
    {
        IQueryProvider _provider;
        IQueryable<T> _inner;

        public ExpandableQuery(IQueryable<T> inner)
        {
            _inner = inner;
            _provider = Activator.CreateInstance<TProvider>().Attach(inner);
        }

        Expression IQueryable.Expression { get { return _inner.Expression; } }
        Type IQueryable.ElementType { get { return typeof(T); } }
        IQueryProvider IQueryable.Provider { get { return _provider; } }
        public IEnumerator<T> GetEnumerator() { return _inner.GetEnumerator(); }
        IEnumerator IEnumerable.GetEnumerator() { return _inner.GetEnumerator(); }
        public override string ToString() { return _inner.ToString(); }
    }

    /// <summary>
    /// The base class for extensible QueryProviders
    /// 
    /// Subclasses can override VisitExpression to extend the expression tree or override the 
    /// IQueryProvider interface for more flexibility
    /// </summary>
    public abstract class AbstractQueryProvider : IQueryProvider
    {
        protected IQueryable _query;

        internal AbstractQueryProvider Attach(IQueryable query)
        {
            _query = query;
            return this;
        }

        public virtual Expression VisitExpression(Expression expression)
        {
            return expression;
        }

        // The following four methods first call VisitExpression to visit the expression tree, then call
        // upon the inner query to do the remaining work.
        IQueryable<TElement> IQueryProvider.CreateQuery<TElement>(Expression expression)
        {
            return (IQueryable<TElement>)Activator.CreateInstance(typeof(ExpandableQuery<,>).MakeGenericType(typeof(TElement), GetType()),
                _query.Provider.CreateQuery<TElement>(VisitExpression(expression)));
        }

        IQueryable IQueryProvider.CreateQuery(Expression expression)
        {
            return _query.Provider.CreateQuery(VisitExpression(expression));
        }

        TResult IQueryProvider.Execute<TResult>(Expression expression)
        {
            return _query.Provider.Execute<TResult>(VisitExpression(expression));
        }

        object IQueryProvider.Execute(Expression expression)
        {
            return _query.Provider.Execute(VisitExpression(expression));
        }
    }


    /// <summary>
    /// Implements expandable expressions
    /// </summary>
    public class ExpandableQueryProvider : AbstractQueryProvider
    {
        public override Expression VisitExpression(Expression expression)
        {
            return base.VisitExpression(expression).Expand();
        }
    }


    /// <summary>
    /// An IQueryable wrapper that allows us to visit the query's expression tree just before LINQ to SQL gets to it.
    /// This is based on the excellent work of Tomas Petricek: http://tomasp.net/blog/linq-expand.aspx
    /// </summary>
    public class ExpandableQuery<T> : IQueryable<T>, IOrderedQueryable<T>, IOrderedQueryable
    {
        ExpandableQueryProvider<T> _provider;
        IQueryable<T> _inner;

        internal IQueryable<T> InnerQuery { get { return _inner; } }			// Original query, that we're wrapping

        internal ExpandableQuery(IQueryable<T> inner)
        {
            _inner = inner;
            _provider = new ExpandableQueryProvider<T>(this);
        }

        Expression IQueryable.Expression { get { return _inner.Expression; } }

        Type IQueryable.ElementType { get { return typeof(T); } }

        IQueryProvider IQueryable.Provider { get { return _provider; } }

        public IEnumerator<T> GetEnumerator() { return _inner.GetEnumerator(); }

        IEnumerator IEnumerable.GetEnumerator() { return _inner.GetEnumerator(); }

        public override string ToString() { return _inner.ToString(); }
    }

    class ExpandableQueryProvider<T> : IQueryProvider
    {
        ExpandableQuery<T> _query;

        internal ExpandableQueryProvider(ExpandableQuery<T> query)
        {
            _query = query;
        }

        // The following four methods first call ExpressionExpander to visit the expression tree, then call
        // upon the inner query to do the remaining work.

        IQueryable<TElement> IQueryProvider.CreateQuery<TElement>(Expression expression)
        {
            return new ExpandableQuery<TElement>(_query.InnerQuery.Provider.CreateQuery<TElement>(expression.Expand()));
        }

        IQueryable IQueryProvider.CreateQuery(Expression expression)
        {
            return _query.InnerQuery.Provider.CreateQuery(expression.Expand());
        }

        TResult IQueryProvider.Execute<TResult>(Expression expression)
        {
            return _query.InnerQuery.Provider.Execute<TResult>(expression.Expand());
        }

        object IQueryProvider.Execute(Expression expression)
        {
            return _query.InnerQuery.Provider.Execute(expression.Expand());
        }
    }
}
